import re

import pandas as pd
from langchain_text_splitters import CharacterTextSplitter

from pasers.xlsx_parser.xlsx_parse import prepare_excel_table


def get_md_tables(table_dic, max_len):
    # 从字典中获取表头和行数据
    table_head = table_dic['table_head']
    row_data = table_dic['row_data']

    # 存储不超过指定长度的行数据
    new_row_data = []

    # 生成合并后的表头Markdown格式并计算长度
    head_md = merge_head(table_head)
    head_md_len = len(head_md)

    # 将表头Markdown拆分成单元格列表，并限制单元格的最大字符长度
    head_cell_ls = head_md.split('\n')[0].split('|')[1:-1]
    head_max_len = int(max_len / 2)

    # 初始化存储超长行数据和超长单元格的列表
    exceed_row_data = []
    exceed_cell_ls = []

    # 如果表头Markdown的长度超过允许的最大长度，对单元格内容进行截断
    if head_md_len > head_max_len:
        for i, head_cell in enumerate(head_cell_ls):
            # 如果某个单元格内容超长，将其截断并记录
            if len(head_cell) > head_max_len:
                exceed_cell_ls.append(head_cell)
                head_cell_ls[i] = head_cell[0:head_max_len]
        # 将截断后的单元格重新组合成Markdown格式的表头
        head_md = to_md_string(head_cell_ls) + '\n' + table_head[-1]

    # 计算每个表头单元格的宽度，包含额外的填充
    head_cell_len_ls = [len(i) + 4 for i in head_cell_ls]

    # 遍历每一行数据，根据最大长度进行截断处理
    for cur_row_md in row_data:
        # 如果当前行的长度加上表头的长度超过最大长度，进行截断处理
        if len(cur_row_md) + head_md_len > max_len:
            # 分割行数据为单元格列表
            row_cell_ls = cur_row_md.split('|')[1:-1]
            # 截断每个单元格的内容，并返回超长的单元格
            row_cell_ls, exceed_cells = truncate_row_text(row_cell_ls, head_cell_len_ls, max_len)
            # 将超长的单元格内容添加到超长单元格列表中
            if exceed_cells:
                exceed_cell_str = ' '.join(exceed_cells)
                exceed_cell_ls.append(exceed_cell_str)
            # 重新生成截断后的Markdown行数据
            cur_row_md = to_md_string(row_cell_ls)
            exceed_row_data.append(cur_row_md)
        else:
            # 如果行数据未超长，直接添加到新行数据列表
            new_row_data.append(cur_row_md)

    # 使用分割后的数据生成子表格
    sub_tables = get_sub_tables(new_row_data, exceed_row_data, head_cell_ls, head_md, max_len)

    # 返回分割后的子表格和超长的单元格列表
    return sub_tables, exceed_cell_ls


def truncate_row_text(row_text, head_cell_len_ls, max_len):
    """
    截断表格行中的单元格文本，使每个单元格及其对应的表头长度之和不超过 max_len。

    :param row_text: 行数据的列表，每个元素是一个单元格的文本内容
    :param head_cell_len_ls: 表头长度列表，包含每个表头单元格的宽度
    :param max_len: 允许的最大长度
    :return:
        rewrite_row_text: 截断后的行数据列表
        exceed_cells: 超过最大长度的原始单元格文本（去除空格）
    """
    extra_len = int(max_len / 10)  # 为Markdown标记符预留的额外长度

    # 存储截断后的行数据
    rewrite_row_text = []

    # 记录超过最大长度限制的单元格内容
    exceed_cells = []

    # 遍历每个单元格及其内容
    for cell_id, cell_text in enumerate(row_text):
        # 判断单元格内容加上表头长度和额外长度是否超过最大允许长度
        if len(cell_text) + head_cell_len_ls[cell_id] + extra_len > max_len:
            # 截断单元格内容，使其符合长度要求
            rewrite_cell_text = cell_text[0:max_len - head_cell_len_ls[cell_id] - extra_len]
            rewrite_row_text.append(rewrite_cell_text)

            # 如果原始单元格内容（去掉多余空格）未被记录过，则添加到超长单元格列表
            if cell_text.strip() not in exceed_cells:
                exceed_cells.append(cell_text.strip())
        else:
            # 若单元格内容不超长，直接添加到行数据中
            rewrite_row_text.append(cell_text)

    # 返回截断后的行数据和超长单元格列表
    return rewrite_row_text, exceed_cells


# 表格切割为子表
def get_sub_tables(row_data, exceed_row_data, head_cell_ls, head_md, max_len):
    """
    将表格数据拆分为多个子表格，确保每个子表格的长度不超过 max_len。

    :param row_data: 普通行数据列表，每个元素是行的Markdown格式字符串
    :param exceed_row_data: 超长行数据列表，每行内容的长度超过最大允许长度
    :param head_md: 表头的Markdown格式字符串
    :param head_cell_ls: 表头数据的单元格列表
    :param max_len: 子表格的最大允许长度
    :return: 拆分后的子表格列表，每个子表格为Markdown格式的字符串
    """

    # 初始化存储子表格的列表
    tables_md_ls = []

    # 处理不超长的行数据
    if row_data:
        # 合并表头和行数据形成表格的Markdown格式
        table_md = head_md + '\n' + '\n'.join(row_data)

        # 如果表格长度在允许范围内，则直接添加
        if len(table_md) < max_len:
            tables_md_ls.extend([table_md])
        else:
            # 若超出长度，将表数据按max_len进行拆分
            data_text = '\n'.join(row_data)
            tables_md_ls.extend(table_split(max_len, data_text, head_md))

    # 处理超长的行数据，逐行检查并进行纵向切割
    if exceed_row_data:
        # 合并表头和超长行数据
        table_md = head_md + '\n' + '\n'.join(exceed_row_data)

        # 解析Markdown表格为标题和数据
        titles, data = md_pase(table_md)

        # 将数据转换为DataFrame格式，便于切割
        df = pd.DataFrame(data, columns=titles)

        # 使用预处理函数将DataFrame切割成多个子表格
        sub_tables_ls, _, _ = prepare_excel_table(df, max_len)

        # 将切割后的子表格加入结果列表
        tables_md_ls.extend(sub_tables_ls)

    # 如果只有表头而无行数据，返回仅包含表头的表格
    if head_cell_ls and not row_data and not exceed_row_data:
        text = '\n'.join(head_cell_ls).strip()
        if len(text) > 0:
            tables_md_ls = [text]

    # 移除子表格中的多余换行和分隔符，确保每个子表格的内容非空
    table_str_ls = []
    for table in tables_md_ls:
        str_table = re.sub('[\n|-]', '', table).strip()
        if len(str_table) > 0:
            table_str_ls.append(table)

    # 返回经过处理的子表格列表
    return table_str_ls


def get_row_text(head_cell_num, row_text):
    """
    使 row_text 中的元素数量和 head_cell_num 中数量相等
    :param head_cell_num:
    :param row_text:
    :return: row_text
    """
    if head_cell_num == len(row_text):
        return row_text
    elif head_cell_num < len(row_text):
        row_text[head_cell_num - 1] = ''.join(row_text[head_cell_num - 1:]).strip()
        return row_text[:head_cell_num]
    else:
        row_text.extend([''] * (head_cell_num - len(row_text)))
        return row_text


def to_md_string(a_ls):
    """
    将 a_ls 列表转换为 markdown 字符串
    :param a_ls:
    :return:
    """
    b = ['|']
    for i in a_ls:
        b.append(str(i).replace('|', '').strip())
        b.append('|')
    return ''.join(b)


def md_pase(md_text):
    md_ls = md_text.split('\n')
    titles = md_ls[0].split('|')[1:-1]
    data = [i.split('|')[1:-1] for i in md_ls[2:]]
    return titles, data


def table_split(max_len, data_text, head_text):
    """
    这里切的是表格的数据部分（data_text），但要保证切割后的数据+表头的长度小于split_len
    :param max_len: 最大切割长度
    :param data_text: 表格的数据部分
    :param head_text: 表头字符串
    :return:
    """
    head_len = len(head_text)
    c_splitter = CharacterTextSplitter(
        separator='\n',
        chunk_size=max_len - head_len,
        chunk_overlap=0
    )
    results = []
    sub_tables = c_splitter.split_text(data_text)
    for sub_table in sub_tables:
        results.append(head_text + '\n' + sub_table)
    return results


def merge_head(table_head):
    """
    table_head为多行表头数据，但表头只能有一行，且表头间数据具有唯一性，故调此函数将多行表头数据合并为一行
    """
    merge_result = [''] * len(table_head[0].split('|'))
    for i, h in enumerate(table_head):
        if i < len(table_head) - 1:
            head_data = h.split('|')
            merge_result = wise_elements_concat(merge_result, head_data)
    return '\n'.join(['|'.join([i.strip() for i in merge_result]), table_head[-1]])


def wise_elements_concat(ls_a, ls_b):
    """
    :param concatenationer: 连接符
    :param ls_a: list, 其中每个元素都是字符串
    :param ls_b: list, 其中每个元素都是字符串
    :return: 将两个列表进行逐元素拼接，要求ls_a,ls_b形状一样
    """
    for i in range(len(ls_a)):
        ls_a[i] = ls_a[i] + ls_b[i]
    return ls_a
